import warnings
warnings.filterwarnings("ignore", category=DeprecationWarning)

import numpy as np
import multiprocessing as mp
from collections import namedtuple
import sys


# This code is for the training / offline phase
ParseParams = namedtuple('ParseParams', 'interval offset contention_frac threshold score')
interval = None
result_x = None
result_y = None
SCORE_MAX = 9999999999999

# 8 kB of random bits
RANDOM_PATTERN = False
pattern = "01011000111100110110000000010000011100001110111101101011010011101011100111001100011000111101001101010101111110100110110110001110001000111100010001101000001111011000111000100100011100000000110011001111001010000100001001110100000111011110111000011110011111000000100100011001110110000110001001011100000110010011111010000001101011010101100011011000110110100010101001011010111010111010010000110001000101100100011010101111100001101100000111010001001000100110101100010100111000001001000000011111111101101100011011101001010001101011100111001110011011110000010110011000101001001001111101110100000011010011110100010000011100110111010100011000101111100010100001101111011000001100001100101111011111110110000111111011110100100000100100010111000001101000010000110100011100010000110000000001011101000100010111001001111001111001001100101001111000111100010100011111110110110000100100111110011110111001011110111110110011110101001001000011100010101010011001110010111110000001001010110000010000010101001110001111011101000010000011101101110001000111011000000101111100100100001110011000101100011100000101001000000001010110001011110000000011011000100111011010000111010111101100001100000010110100010001100000111100011101010011000101100110110110110111010000110111010001011001111101001110000010100010110000110001101111101111001000100010010101100100101110100000101110101100001001101011010100011100111001100010110000101100111000000101100001010000001000111001000001110110101110111000001110000000011011011100010111001010010101100100101011001010011110111101000100100000010011001100010110001001010011011100010000010100001100101110011101011000011100101100000111111101011100111110101011101011111011010001101110110110110100001101100101101110010011011100010010000110001101110100011000001011111010000101010101010011011100001000101111011110100011100001010101010100110110001110010110010011101000001100111100110101010001110001010101001110100101111011100101010100101011000100100111010111000110001100010100010100100100110111101010100011000100110000111100111011000000001010011011010010000111101110000010101100111001010000011101011001100111010111110101101110110010011010100001010100011110011000000101111111010010100000000000110101010110111011011110101101000010010110001000000001101101111011000001011011011001101001010110011100011110011000000101010101101000111001000001011011011110001110001101101010010000011011110101000100000101001101100110010110110100101100111110000100111110011110011101100001001110011100010110001000110110011001011101001000101110010111110010111101000101011111111111010110010001011100111101101111001110101101000011101011000110110110000111001111101101111011100001111010110101010100110010011111000000000011100010101001010011111111011001001011110001111000110001111010110000100011010000000001010110111110111011111111011011011011100010000011101001101011110101100100101001011111110001011111000001111110000001001101001100100001001001011000100100111100010010001111111111001001110000110111110000000000111010001110010111100011101100110010100100111111101000100000011000000010111110000111001000111100111010111101101101000001000011010111000100001110110000010100000000111101111001100101001011001000100100111001010100100111010000101101011000100011100110000101101110101100000110010110110010001101010110111001101101001101001000000000110010000000101110111000101010001010010001111101101101001110101110011000011111000110011110000101110110101100001110001000111110000100100011101001100111011000101010001101100000011011100010001000000000110000100010011000001000001100010100011001011100110011111000000110111100000001000101100001000010101011010011000010111101100101011001110101110101111000101100110010001111010111000000101001010010101011010011110101111100010100000111101100101001010111011000111100010000001100100111100101100111111010110011111100101100101100101110010001101110100001110101000110001001100001110001001100110010010111011101110011000001111011010010101011101110000010001000010100110000001101010011000100111110100100101010000100001011011100110011011101100110110101011111101100000010000010011101110010110101111011000001010101111001000100110111011100101101000100111001100111101011101010011000011000001110100100101011000101010011000010101010000111100011000010000101110011000001101100001110001001010101010011000100000000101011111000011000000001010001111001001010010010110000010010010001111100101101010011000011100000011111000010111100111100101001101100010110110000010010001001101000010100111001000010000100011101011010000101011011111001011100010001010010111111000011000101000010001000100101000001100101100001010001111111010100010000100001011100100100101100100000110000101000010101010000011100001001011010101101111101001110100010010011101001000110000101100011100001001101100010100101111111101110100110011011101010000101101010110000011000010100101001111011110101110010001100111001111001000001000001000001000010101100011010111111011100000111011110000011010101011010111100001111001011010011110000010100011101110100000001011010001000100001010110100100000111000011000111011101000011111000101000101111101100001101110011010001111111001011111000000000110010111000011000111011110010111111110111100000110010110101011100011010111100011111100000101101110010010110100110111100101000111001101011111100101011111101000000100110111001110110110011101001001000001100001001000100110000101000001000110111110101000100111100111000011011011000110011010101010000100000100001011100111010101010100110001001000101111001000111100010100011011110101010000111001110100000011001010100001011111110110100111100100000100011000111011101011001010100001101010111111110100001010011001000000101111011101110000001010001101111011011000110101000000011101010111101111011101100100001110001001000100100011011010110010010100001100000010100110100001011001001001001100000011101110101000011101111110111101010100011010001011000110100111110110110110001010111110110000000111010010010010110101010010100111110111100011110101100110011111001001000000110101010001010010111011111000010000010100111000101000001011110010100001100101100101001100000011000110101000000010011011011101111110001111110110111000110110110001001101011100010011001000011010010010010010111101100101011001101110000001100001111000100001001000110100001010110011100111001101101001111101111101010110011100101111000001100101000111001101111100011001000010000101111010111100110101100011001111110001001111000100101001001011001001010000100011000110011111101000000110110001111100111111100100001100010010111001010001100000001000010100110011100101111101010011000101101110011010110100010100111100000010011100111001000001010100011101000001000011110011010001110010001111000101000011010101100110000111100101110000110001100111010011110011100101001000111111010100000100001010001100100011100000101010101001101000101111001111111111001100100011001010100010000111111010110010100011001110100111001010110000100010111001100101101011111111000101000111110001111100110001011101001010100111101111000001000001000000100010011011010110110000010100000111111100011000111100111101111011100011001011111001100001111101001010010111001010000000101010100001001101101110110111100111110011000010110010111001110111111101001000001100110100110111110010010111110101000111100010111010111011001111110100001000001100101100000100001011101111011000010100001100101100110010111111000101101001100001000101001101011101011111011100101111110111111111111101011110111011010100101110000100011110010111000000001001110100110110110000001000100101110011010011100111011111111110000110111010111011010000110000111010100110010000101000100110011100111001101001101100101011000111101101010110101000000101000111111011101101111010000111000101101100110010111001001001000000101110011001100110001010001100000000010000011110100010010010101111000011110011010000010010110011110101001000101101000110000100100101111111001100111110011100100110010011001011011110011011101010110111110011111001011001011010001011011111010000110000100100011001110111011010111110110101001110010100010100100110111001001100100100111010010110100"
patternlen = len(pattern)

# The first 100 intervals are discarded
# The following 8000 intervals are the training set (offline phase)
# The following 8000 intervals are the testing set (online phase)
# NOTE: if RANDOM_PATTERN is True, make sure that the number of train and test intervals is the same as patternlen!
discard_intervals = 100
train_intervals_no = 8000
test_intervals_no = 8000

discard_intv_start = 0
discard_intv_end = discard_intervals
train_intv_start = discard_intervals
train_intv_end = train_intv_start + train_intervals_no
test_intv_start = discard_intervals + train_intervals_no
test_intv_end = test_intv_start + test_intervals_no


def read_from_file(filename):
    result_x = []
    result_y = []

    # Results of all runs are all in the same file
    with open(filename) as f:
        for line in f:
            x, y = line.strip().split()
            result_x.append(int(x))
            result_y.append(int(y))

    return result_x, result_y


def diff_letters(a, b):
    return sum(a[i] != b[i] for i in range(len(a)))


def parse_intervals_into_bits(intervals, threshold, min_contention_frac):

    # Parse the intervals into bits
    # If more than $(min_contention_frac)% of the measurements in an interval
    # are greater than the threshold, then we classify
    # the interval as a 1.
    result = ""
    for _, samples in intervals.items():
        contention = 0
        for sample in samples:
            if sample >= threshold:
                contention += 1
        if contention > min_contention_frac * len(samples):
            result += "1"
        else:
            result += "0"

    return result


def per_offset_worker(parse_params):

    # Parse trace into intervals
    offset = parse_params.offset
    interval = parse_params.interval

    # Prepare training set taking into account the offset for this worker
    cur_interval_no = 0
    train_intervals = {}
    for i in range(len(result_x)):
        x = result_x[i]
        y = result_y[i]
        if ((x + offset) > interval * (cur_interval_no + 1)):
            cur_interval_no += 1

        if (discard_intv_start <= cur_interval_no < discard_intv_end):
            continue
        if (train_intv_start <= cur_interval_no < train_intv_end):
            train_intervals.setdefault(cur_interval_no, []).append(y)
        else:
            break

    # We use the training set to find the best threshold
    # for this interval (offline) and the remaining parsed data
    # to get the error rate (online). We try different thresholds
    # because, depending on the CC interval, the best threshold
    # to use is different (e.g., with larger intervals we can use a
    # lower threshold than with shorter intervals).
    # We try different fractions of contention samples observed (e.g. 10% of
    # the samples must show contention for the bit to be counted as a 1)
    thresholds = range(4000000, 4300000, 10000)  # FIXME: adjust these thresholds for your CPU
    contention_fracs = range(0, 100, 3)  # test from 0.03 to 1.00

    # Save best scores
    best_contention_frac = None
    best_threshold = None
    best_score = SCORE_MAX

    # Test various contention fractions
    for min_contention_frac in contention_fracs:
        for threshold in thresholds:

            # Parse the intervals into bits
            result = parse_intervals_into_bits(train_intervals, threshold, min_contention_frac / 100)

            if RANDOM_PATTERN:
                score = patternlen

                # Find the best offset using the first patternlen intervals
                nppattern = np.array([int(i) for i in pattern])
                npresult = np.array([int(i) for i in result])
                for i in range(patternlen):
                    testing = np.roll(nppattern, i)
                    newscore = np.count_nonzero(testing != npresult)

                    if (newscore < score):
                        score = newscore
            else:
                # Compare these bits with the ground truth
                # The ground truth is either 0101... or 1010...
                candidate_1 = "01" * (len(result) // 2)
                candidate_2 = "10" * (len(result) // 2)

                # Get the number of bit flips between the
                # decoded stream and the (correct) ground truth
                score_1 = diff_letters(result, candidate_1)
                score_2 = diff_letters(result, candidate_2)
                score = min(score_1, score_2)

            # Pick the best score (lowest error)
            if (score < best_score):
                best_threshold = threshold
                best_contention_frac = min_contention_frac / 100
                best_score = score

    return ParseParams(interval, offset, best_contention_frac, best_threshold, best_score)


def main():
    global interval, result_x, result_y

    # Check args
    assert len(sys.argv) == 3, "Specify the files with the results as argument, and the interval"
    result_x, result_y = read_from_file(sys.argv[1])

    interval = int(sys.argv[2])

    # 1 milliseconds is 3 million cycles on our processor
    # FIXME: adjust this with correct base frequency
    samples_per_interval = interval // 3000000

    # Train using different offsets
    pool = mp.Pool(processes=8)
    offsets = range(0, interval // 2, interval // samples_per_interval)
    params = [ParseParams(interval, o, None, None, None) for o in offsets]
    best_params_per_offset = pool.map(per_offset_worker, params)

    # Select the best parameters
    best_params = min(best_params_per_offset, key=lambda x: x.score)
    best_offset = best_params.offset
    best_threshold = best_params.threshold
    best_contention_frac = best_params.contention_frac

    # Generate test set
    cur_interval_no = 0
    test_intervals = {}
    for i in range(len(result_x)):
        x = result_x[i]
        y = result_y[i]
        if ((x + best_offset) > interval * (cur_interval_no + 1)):
            cur_interval_no += 1

        if (test_intv_start <= cur_interval_no < test_intv_end):
            test_intervals.setdefault(cur_interval_no, []).append(y)
        elif cur_interval_no > test_intv_end:
            break
        else:
            continue

    if len(test_intervals) < test_intervals_no:
        print("Not enough test intervals")
        exit(-1)

    print('Best params: {}'.format(best_params))

    # Parse the intervals into bits
    result = parse_intervals_into_bits(test_intervals, best_threshold, best_contention_frac)

    # Check resulting number of errors
    if RANDOM_PATTERN:
        score = SCORE_MAX

        # Find the best offset using the first patternlen intervals
        nppattern = np.array([int(i) for i in pattern])
        npresult = np.array([int(i) for i in result])
        for i in range(patternlen):
            testing = np.roll(nppattern, i)
            newscore = np.count_nonzero(testing != npresult)

            if (newscore < score):
                score = newscore

    else:
        # Compare these bits with the ground truth
        # The ground truth is either 0101... or 1010...
        candidate_1 = "01" * ((len(result) // 2) + 1)
        candidate_2 = "10" * ((len(result) // 2) + 1)

        # Print the number of bit flips between the
        # decoded stream and the (correct) ground truth
        score_1 = diff_letters(result, candidate_1[:len(result)])
        score_2 = diff_letters(result, candidate_2[:len(result)])
        score = min(score_1, score_2)

    print(score, "out of", test_intervals_no, "->", score / test_intervals_no)


if __name__ == "__main__":
    main()
